#!/usr/bin/env python3
"""
视频关键帧提取脚本
使用 OpenCV 从视频中提取关键帧
"""

import cv2
import os
import numpy as np
from pathlib import Path
import argparse

def calculate_frame_difference(frame1, frame2):
    """计算两帧之间的差异"""
    # 转换为灰度图
    gray1 = cv2.cvtColor(frame1, cv2.COLOR_BGR2GRAY)
    gray2 = cv2.cvtColor(frame2, cv2.COLOR_BGR2GRAY)
    
    # 计算结构相似性
    diff = cv2.absdiff(gray1, gray2)
    return np.mean(diff)

def extract_keyframes(video_path, output_dir, method='time', interval=1.0, threshold=30.0, max_frames=20):
    """
    从视频中提取关键帧
    
    参数:
    video_path: 视频文件路径
    output_dir: 输出目录
    method: 提取方法 ('time' 按时间间隔, 'difference' 按帧差异)
    interval: 时间间隔(秒) 或 帧差异阈值
    threshold: 帧差异阈值
    max_frames: 最大提取帧数
    """
    
    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)
    
    # 打开视频文件
    cap = cv2.VideoCapture(video_path)
    if not cap.isOpened():
        print(f"❌ 无法打开视频文件: {video_path}")
        return False
    
    # 获取视频信息
    fps = cap.get(cv2.CAP_PROP_FPS)
    total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
    duration = total_frames / fps if fps > 0 else 0
    
    print(f"🎬 视频信息:")
    print(f"   FPS: {fps:.2f}")
    print(f"   总帧数: {total_frames}")
    print(f"   时长: {duration:.2f} 秒")
    print(f"   提取方法: {method}")
    
    keyframes = []
    frame_count = 0
    prev_frame = None
    
    if method == 'time':
        # 按时间间隔提取
        frame_interval = int(fps * interval)
        print(f"   时间间隔: {interval} 秒 ({frame_interval} 帧)")
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
                
            if frame_count % frame_interval == 0 and len(keyframes) < max_frames:
                timestamp = frame_count / fps
                keyframes.append((frame_count, frame, timestamp))
                print(f"   提取帧 {frame_count} (时间: {timestamp:.2f}s)")
            
            frame_count += 1
            
    elif method == 'difference':
        # 按帧差异提取
        print(f"   差异阈值: {threshold}")
        
        while True:
            ret, frame = cap.read()
            if not ret:
                break
            
            if prev_frame is not None and len(keyframes) < max_frames:
                diff = calculate_frame_difference(prev_frame, frame)
                
                if diff > threshold:
                    timestamp = frame_count / fps
                    keyframes.append((frame_count, frame.copy(), timestamp))
                    print(f"   提取帧 {frame_count} (时间: {timestamp:.2f}s, 差异: {diff:.2f})")
            
            # 保存第一帧
            elif prev_frame is None and len(keyframes) < max_frames:
                timestamp = frame_count / fps
                keyframes.append((frame_count, frame.copy(), timestamp))
                print(f"   提取帧 {frame_count} (首帧)")
            
            prev_frame = frame.copy()
            frame_count += 1
    
    # 如果没有提取到足够的关键帧，补充一些均匀分布的帧
    if len(keyframes) < 5:
        print("   关键帧不足，补充均匀分布的帧...")
        cap.set(cv2.CAP_PROP_POS_FRAMES, 0)
        step = max(1, total_frames // 10)  # 每10%提取一帧
        
        for i in range(0, total_frames, step):
            if len(keyframes) >= max_frames:
                break
                
            cap.set(cv2.CAP_PROP_POS_FRAMES, i)
            ret, frame = cap.read()
            if ret:
                timestamp = i / fps
                # 检查是否已经有相似的帧
                is_duplicate = False
                for existing_frame_num, existing_frame, _ in keyframes:
                    if abs(existing_frame_num - i) < fps:  # 1秒内的帧认为重复
                        is_duplicate = True
                        break
                
                if not is_duplicate:
                    keyframes.append((i, frame, timestamp))
                    print(f"   补充帧 {i} (时间: {timestamp:.2f}s)")
    
    cap.release()
    
    # 保存关键帧
    video_name = Path(video_path).stem
    saved_count = 0
    
    for i, (frame_num, frame, timestamp) in enumerate(keyframes):
        # 生成文件名
        filename = f"{video_name}_frame_{frame_num:06d}_{timestamp:.2f}s.jpg"
        output_path = os.path.join(output_dir, filename)
        
        # 保存图片
        success = cv2.imwrite(output_path, frame)
        if success:
            saved_count += 1
            print(f"✅ 保存关键帧: {filename}")
        else:
            print(f"❌ 保存失败: {filename}")
    
    print(f"\n🎯 提取完成!")
    print(f"   总共提取了 {len(keyframes)} 个关键帧")
    print(f"   成功保存了 {saved_count} 个图片文件")
    print(f"   保存位置: {output_dir}")
    
    return True

def main():
    parser = argparse.ArgumentParser(description='提取视频关键帧')
    parser.add_argument('--video', '-v', default='/home/hxz/cc/miniCompose/tool/movie/2024530431.mp4', help='视频文件路径')
    parser.add_argument('--output', '-o', default='/home/hxz/cc/miniCompose/tool/movie/frames', help='输出目录')
    parser.add_argument('--method', '-m', choices=['time', 'difference'], default='time', help='提取方法')
    parser.add_argument('--interval', '-i', type=float, default=2.0, help='时间间隔(秒)')
    parser.add_argument('--threshold', '-t', type=float, default=30.0, help='帧差异阈值')
    parser.add_argument('--max-frames', '-n', type=int, default=15, help='最大提取帧数')
    
    args = parser.parse_args()
    
    if not os.path.exists(args.video):
        print(f"❌ 视频文件不存在: {args.video}")
        return
    
    print(f"🎬 开始提取关键帧...")
    print(f"   输入视频: {args.video}")
    print(f"   输出目录: {args.output}")
    
    success = extract_keyframes(
        video_path=args.video,
        output_dir=args.output,
        method=args.method,
        interval=args.interval,
        threshold=args.threshold,
        max_frames=args.max_frames
    )
    
    if success:
        print("🎉 关键帧提取任务完成！")
    else:
        print("❌ 关键帧提取失败！")

if __name__ == "__main__":
    main()